﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace Alignmentizer
{
	class Matrix
	{

		private cell[,] _matrix;
		private string _seq1, _seq2;
		private int _match;
		private int _mismatch;
		private int _gap;
		private Dictionary<string, int> _scoringMatrix;

		public enum Direction {up, left, diagonal, none};

		public struct cell
		{
			public int value;
			public Direction dir;
		}

		/// <summary>
		/// Initializes the matrix
		/// </summary>
		/// <param name="seq1"></param>
		/// <param name="seq2"></param>
		/// <param name="match"></param>
		/// <param name="mismatch"></param>
		/// <param name="gap"></param>
		/// <param name="global">True for global, false for local</param>
		public Matrix(string seq1, string seq2, int match = 5, int mismatch = -3, int gap = -4, bool global = true, string scoringMatrixPath = "")
		{
			_match = match;
			_mismatch = mismatch;
			_gap = gap;
			_seq1 = "-" + seq1;	//rows
			_seq2 = "-" + seq2;	//columns
			_matrix = new cell[_seq2.Length, _seq1.Length];
			if (scoringMatrixPath != "")
				_scoringMatrix = HelperFunctions.ImportScoringMatrix(scoringMatrixPath);

			if (global)
			{
				//initialize first row and column of matrix
				for (int i = 0; i < _seq1.Length; i++)
					_matrix[0, i].value = _gap * i;
			
				for (int i = 0; i < _seq2.Length; i++)
					_matrix[i, 0].value = _gap * i;
			}
			
			else
			{
				for (int i = 0; i < _seq1.Length; i++)
				{
					_matrix[0, i].value = 0;
					_matrix[0, i].dir = Direction.none;
				}

				for (int i = 0; i < _seq2.Length; i++)
				{
					_matrix[i, 0].value = 0;
					_matrix[i, 0].dir = Direction.none;
				}
			}

			MatrixFill(global);

		}

		/// <summary>
		/// Fills up the matrix according to the rules set forth by the GlobalAlignmentMax function
		/// </summary>
		public void MatrixFill(bool global)
		{
			int iMax=0, jMax=0, maxValue=0;
			for (int i = 1; i < _seq2.Length; i++)
			{
				for (int j = 1; j < _seq1.Length; j++)
				{
					if (global)
						_matrix[i, j] = GlobalAlignmentMax(i, j);
					else
						_matrix[i, j] = LocalAlignmentMax(i, j);
						
					//Keep track of the maximum value in the matrix to be used as the start of the traceback
					if (_matrix[i, j].value > maxValue)
					{
						maxValue = _matrix[i, j].value;
						iMax = i;
						jMax = j;
					}
				}
			}

			if (global)
				TraceBack();
			else
				TraceBack(iMax, jMax);

			Console.WriteLine("\nMax alignment score: {0}", _matrix[iMax, jMax].value);

		}

		/// <summary>
		/// Prints the matrix to the console
		/// </summary>
		public void PrintMatrix()
		{
			for(int i = 0; i < _seq2.Length; i++)
			{
				Console.Write("{0}", _matrix[i, 0].value);
				for(int j = 1; j < _seq1.Length; j++)
					Console.Write("\t{0}", _matrix[i, j].value);
				Console.WriteLine();
			}

		}

		/// <summary>
		/// Determines the alignment score for a given cell using the Needleman-Wunsch algorithm
		/// </summary>
		/// <param name="i"></param>
		/// <param name="j"></param>
		/// <returns></returns>
		private cell GlobalAlignmentMax(int i, int j)
		{
			cell result = new cell();
			int diagScore;

			if(_seq1[j] == _seq2[i])
				diagScore = _match;
			else
				diagScore = _mismatch;

			int diagFinalScore = _matrix[i-1,j-1].value + diagScore;
			int leftScore = _matrix[i, j-1].value + _gap;
			int upScore = _matrix[i-1, j].value + _gap;

			if (diagFinalScore < leftScore)
				if (upScore < leftScore)
				{
					//left score wins case
					result.value = leftScore;
					result.dir = Direction.left;
				}
				else
				{
					//up score wins case
					result.value = upScore;
					result.dir = Direction.up;
				}
			else if (diagFinalScore < upScore)
				{
					//up score wins case
					result.value = upScore;
					result.dir = Direction.up;
				}
			else
			{
				//diagonal score wins case
				result.value = diagFinalScore;
				result.dir = Direction.diagonal;
			}
			

			return result;
		}

		/// <summary>
		/// Determines the value of a given cell according to the Smith-Waterman algorithm
		/// </summary>
		/// <param name="i"></param>
		/// <param name="j"></param>
		/// <returns></returns>
		private cell LocalAlignmentMax(int i, int j)
		{
			cell result = new cell();
			int diagScore;

			if (_seq1[j] == _seq2[i])
				diagScore = _match;
			else
				diagScore = _mismatch;

			string pair = "";
			pair += _seq1[j];
			pair += _seq2[i];

			if (_scoringMatrix.ContainsKey(pair))
				diagScore = _scoringMatrix[pair];
			else
			{
				string reversePair = "";
				reversePair += pair[1];
				reversePair += pair[0];
				diagScore = _scoringMatrix[reversePair];
			}

			int diagFinalScore = _matrix[i - 1, j - 1].value + diagScore;
			int leftScore = _matrix[i, j - 1].value + _gap;
			int upScore = _matrix[i - 1, j].value + _gap;

			if (diagFinalScore < leftScore)
				if (upScore < leftScore)
				{
					//left score wins case
					result.value = leftScore;
					result.dir = Direction.left;
				}
				else
				{
					//up score wins case
					result.value = upScore;
					result.dir = Direction.up;
				}
			else if (diagFinalScore < upScore)
			{
				//up score wins case
				result.value = upScore;
				result.dir = Direction.up;
			}
			else
			{
				//diagonal score wins case
				result.value = diagFinalScore;
				result.dir = Direction.diagonal;
			}

			//the difference between local and global here: if the value is less than zero then set it equal to zero and set direction to none
			if (result.value <= 0)
			{
				result.value = 0;
				result.dir = Direction.none;
			}

			return result;
		}

		/// <summary>
		/// Traces back through the matrix solving everything in its path
		/// </summary>
		public void TraceBack()
		{
			TraceBack(_seq2.Length - 1, _seq1.Length - 1);
		}

		/// <summary>
		/// Traces back through the matrix solving everything in its path
		/// </summary>
		/// <returns></returns>
		public void TraceBack(int istart, int jstart)
		{
			string[] result = new  string[2];
			result[0] = "";
			result[1] = "";
			int writeoutLength = 80;

			result = TraceBack(istart, jstart, result);

			//Write out the aligned sequences writeoutLength characters at a time with the matches pointed out
			int index = 0;
			while (index != result[0].Length)
			{
				if (writeoutLength > result[0].Length - index)
					writeoutLength = (result[0].Length - index);

				for (int i = 0; i < writeoutLength; i++)
					Console.Write(result[1][index + i]);
				Console.WriteLine();

				for (int i = 0; i < writeoutLength; i++)
				{
					if (result[0][index+i] == result[1][index+i])
						Console.Write("|");
					else
						Console.Write(" ");
				}
				Console.WriteLine();

				for (int i = 0; i < writeoutLength; i++)
					//index is incremented here so the program will print the next block of characters
					Console.Write(result[0][index++]); 
				
				Console.WriteLine();
				Console.WriteLine();
			}

		}

		/// <summary>
		/// Recrursively trace back through the matrix doing shit
		/// </summary>
		/// <param name="i"></param>
		/// <param name="j"></param>
		/// <param name="str1"></param>
		/// <param name="str2"></param>
		/// <returns></returns>
		private string[] TraceBack(int i, int j, string[] sequences)
		{
			string[] result = new string[2];
			//case done
			if (i == 0 && j == 0)
				return sequences;
			
			else if (_matrix[i, j].dir == Direction.none)
				return sequences;
			
			//If i or j are zero then you can't decrease them any more. No need to do any math
			else if (i == 0)
			{
				sequences[0] = "-" + sequences[0];
				sequences[1] = _seq1[j] + sequences[1];
				return TraceBack(i, j - 1, sequences);
			}
			else if (j == 0)
			{
				sequences[0] = _seq2[i] + sequences[0];
				sequences[1] = "-" + sequences[1];
				return TraceBack(i - 1, j, sequences);
			}
			//case up
			else if (_matrix[i, j].dir == Direction.up)
			{
				sequences[0] = _seq2[i] + sequences[0];
				sequences[1] = "-" + sequences[1];
				return TraceBack(i - 1, j, sequences);
			}

			//case left
			else if (_matrix[i, j].dir == Direction.left)
			{
				sequences[0] = "-" + sequences[0];
				sequences[1] = _seq1[j] + sequences[1];
				return TraceBack(i, j - 1, sequences);
			}

			//case diagonal
			else
			{
				sequences[0] = _seq2[i] + sequences[0];
				sequences[1] = _seq1[j] + sequences[1];
				return TraceBack(i - 1, j - 1, sequences);
			}
		}

	}
}
